from stock import Stock
from strategy import Strategy
from log import Log
from data import FileIO

class Backtest(Stock):
    def __init__(self, data, Strategy, start_day, end_day, origin_total_value, period, cycle):
        Stock.__init__(self, data)
        self.strategy = Strategy
        self.start_idx = self.dates.index(start_day) # 开始日期索引值
        self.end_idx = self.dates.index(end_day) # 结束日期索引值
        self.origin_total_value = origin_total_value # 初始资金
        self.total_value = origin_total_value # 目前份额
        self.total_cash = origin_total_value # 目前资金
        self.hold_stock = {} # 目前持仓股票
        self.period = period # 回测时间段
        self.cycle = cycle # 调仓周期
        self.log = Log("data_log.csv")
        self.base_log = Log("base_log.csv")

    # 买入
    def buy(self, id, date, target_value, flag):
        # 买入的手数（1手 = 100股）
        hands = target_value / self.close_price(id, date) //100
        if hands > 0:
            self.hold_stock[id] = [hands, self.close_price(id, date)]
            # 手续费万分之3
            cost = hands * 100 * self.close_price(id, date) * 0.0003
            self.total_value -= cost
            self.total_cash -= (hands * 100 * self.close_price(id, date) + cost)
            operation_content = date + ": 买入股票 " + id + " 共 " + str(int(hands)) + " 手"
            if flag == "normal":
                FileIO().write_txt(operation_content)
            print("买入股票" , id , "共" , str(int(hands)) , "手")

    # 卖出
    def sell(self, id, date):
        self.total_value += (self.close_price(id, date) - self.hold_stock[id][1]) * self.hold_stock[id][0] * 100
        self.total_cash +=self.close_price(id, date) * self.hold_stock[id][0] * 100
        # 手续费万分之3
        cost = self.hold_stock[id][0] * 100 * self.close_price(id, date) * 0.0003
        self.total_value -= cost
        self.total_cash -= cost
        operation_content = date + ": 卖出股票 " + id + " 共 " + str(int(self.hold_stock[id][0])) + " 手"
        FileIO().write_txt(operation_content)
        print("卖出股票" , id ,"共", str(int(self.hold_stock[id][0])), "手")
        del self.hold_stock[id]

    #交易
    def trade(self, date, buy_list, flag):
        print(date)
        # 买入列表的长度为0时，设其为已持有的股票列表（从而减少卖出的手续费开销）
        if len(buy_list) == 0:
            buy_list = list(self.hold_stock.keys())
        # 卖出目前持有股票
        for id in list(self.hold_stock.keys()):
            # 要买入的股票已经持有就不再卖出
            if id not in buy_list:
                self.sell(id, date)
            else:
                # 更新当前份额和所持股票的收盘价
                self.total_value += (self.close_price(id, date) - self.hold_stock[id][1]) * self.hold_stock[id][0] * 100
                self.hold_stock[id][1] = self.close_price(id, date)

        if len(buy_list) > 0:
            # 要买的每只股票的份额
            per_stock = self.total_cash / len(buy_list)
            # 买入股票
            for id in buy_list:
                # 要买入的股票若已经持有就不再买入
                if id not in list(self.hold_stock.keys()):
                    self.buy(id, date, per_stock, flag)
        print("目前持仓总市值：", round(self.total_value, 2))
        # 计算收益率
        rate = (self.total_value - self.origin_total_value) / self.origin_total_value
        rate = format(rate, '.2%')
        print('收益率：', rate)
        hold = {}
        for id in list(self.hold_stock.keys()):
            hold[id] = int(self.hold_stock[id][0])
        if flag == "normal":
            # total_value和rate均保留两位小数
            self.log.record_log(date, round(self.total_value, 2), rate, hold)
        else:
            self.base_log.record_log(date, round(self.total_value, 2), rate, hold)

    def run(self):
        for i in range(self.start_idx, self.end_idx + 1, self.cycle):
            # 从策略中获得买入股票的列表
            # buy_list = self.strategy.choose(self.dates[i], 2, self.period)
            buy_list = self.strategy.choose2(self.dates[i], self.period)
            # 进行交易
            self.trade(self.dates[i], buy_list, "normal")

        # 计算基准收益
        self.total_value = self.origin_total_value
        self.total_cash = self.origin_total_value
        self.hold_stock = {}
        for i in range(self.start_idx, self.end_idx + 1, self.cycle):
            # 进行交易
            self.trade(self.dates[i], ["000001.XSHE"], "base")


